4
https://raw.githubusercontent.com/succlz123/StarDriver-APT/master/stardriver-apt-processor/src/main/java/org/succlz123/stardriver/StarDriverProcessor.java
package org.succlz123.stardriver;

import com.google.auto.service.AutoService;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.ParameterizedTypeName;
import com.squareup.javapoet.TypeSpec;
import com.squareup.javapoet.TypeVariableName;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Stack;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Messager;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.Processor;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.PackageElement;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import javax.tools.Diagnostic;

@AutoService(Processor.class)
public class StarDriverProcessor extends AbstractProcessor {
    private static final String TAG = "Star Driver";
    private static final String PACKAGE_NAME = "org.succlz123.stardriver";
    private Messager mMessager;
    private Elements mElementUtils;
    private Map<String, ProcessClassNode> mNodeHashMap = new HashMap<>();

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        mMessager = processingEnv.getMessager();
        mElementUtils = processingEnv.getElementUtils();
    }

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        HashSet<String> supportTypes = new LinkedHashSet<>();
        supportTypes.add(StarDriverInit.class.getCanonicalName());
        return supportTypes;
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    @Override
    public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnvironment) {
        long startTime = System.currentTimeMillis();
        Set<? extends Element> elements = roundEnvironment.getElementsAnnotatedWith(StarDriverInit.class);
        if (elements != null && !elements.isEmpty()) {
            mMessager.printMessage(Diagnostic.Kind.NOTE, TAG + " -> processing");
            mNodeHashMap.clear();
            for (Element element : elements) {
                String fullName = element.asType().toString();
                ProcessClassNode proxy = mNodeHashMap.get(fullName);
                if (proxy == null) {
                    proxy = new ProcessClassNode(new ClassParameter(mElementUtils, element));
                    mNodeHashMap.put(fullName, proxy);
                }
            }
            try {
                JavaFile javaFile = JavaFile.builder(PACKAGE_NAME, generateManagerCode()).build();
                javaFile.writeTo(processingEnv.getFiler());
            } catch (Exception e) {
                mMessager.printMessage(Diagnostic.Kind.ERROR, TAG + " -> process error " + e);
            }
            double cost = (System.currentTimeMillis() - startTime) / 1000d;
            mMessager.printMessage(Diagnostic.Kind.NOTE, TAG + " -> process finish cost time: " + cost + "s");
        }
        return true;
    }

    private TypeSpec generateManagerCode() {
        return TypeSpec.classBuilder("StarDriverManager")
                .addModifiers(Modifier.PUBLIC)
                .addJavadoc("This Java file is automatically generated by " + TAG + ", PLEASE DO NOT EDIT!")
                .addMethod(generateMethods())
                .build();
    }

    private MethodSpec generateMethods() {
        ClassName host = ClassName.get("android.app", "Application");
        String packageName = this.getClass().getPackage().getName();
        ClassName result = ClassName.get(packageName, "StarDriverStatistics");
        ClassName arrayList = ClassName.get(ArrayList.class);
        ParameterizedTypeName returnType = ParameterizedTypeName.get(arrayList, result);
        MethodSpec.Builder methodBuilder = MethodSpec.methodBuilder("inject")
                .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                .returns(returnType)
                .addParameter(host, "context");
        methodBuilder.beginControlFlow("if (context == null)")
                .addStatement("throw new IllegalStateException(\"The application is null\")")
                .endControlFlow();
        methodBuilder.addStatement("$T<$T> resultList = new $T<>()", ArrayList.class, result, ArrayList.class);

        Set<Map.Entry<String, ProcessClassNode>> entries = mNodeHashMap.entrySet();
        for (Map.Entry<String, ProcessClassNode> entry : entries) {
            ProcessClassNode currentNode = entry.getValue();
            ClassParameter currentParameter = entry.getValue().value;
            // https://stackoverflow.com/questions/7687829/java-6-annotation-processing-getting-a-class-from-an-annotation/52793839#52793839
            try {
                currentParameter.starDriverInit.after();
            } catch (javax.lang.model.type.MirroredTypesException mte) {
                List<? extends TypeMirror> typeMirrors = mte.getTypeMirrors();
                for (TypeMirror typeMirror : typeMirrors) {
                    ProcessClassNode node = mNodeHashMap.get(typeMirror.toString());
                    currentNode.count++;
                    node.next.add(currentNode);
                }
            }
            try {
                currentParameter.starDriverInit.before();
            } catch (javax.lang.model.type.MirroredTypesException mte) {
                List<? extends TypeMirror> typeMirrors = mte.getTypeMirrors();
                for (TypeMirror typeMirror : typeMirrors) {
                    ProcessClassNode node = mNodeHashMap.get(typeMirror.toString());
                    node.count++;
                    currentNode.next.add(node);
                }
            }
        }
        List<ClassParameter> list = new ArrayList<>(entries.size());
        Stack<ProcessClassNode> stack = new Stack<>();
        for (Map.Entry<String, ProcessClassNode> entry : entries) {
            if (entry.getValue().count == 0) {
                stack.push(entry.getValue());
            }
        }
        while (!stack.isEmpty()) {
            ProcessClassNode processClassNode = stack.pop();
            if (processClassNode == null) {
                break;
            }
            for (ProcessClassNode node : processClassNode.next) {
                node.count--;
                if (node.count == 0) {
                    stack.push(node);
                }
            }
            list.add(processClassNode.value);
        }

        if (list.size() != entries.size()) {
            for (ClassParameter classParameter : list) {
                mNodeHashMap.remove(classParameter.classFullName);
            }
            StringBuilder sb = new StringBuilder("The init tasks is Cycle, please check it carefully -> \n");
            for (Map.Entry<String, ProcessClassNode> entry : entries) {
                sb.append(entry.getValue().value.classFullName);
                sb.append("\n");
            }
            throw new IllegalStateException(sb.toString());
        }

        for (ClassParameter parameter : list) {
            methodBuilder.addJavadoc("\"" + parameter.simpleName + "\"" + " - >");
        }

        methodBuilder.addStatement("long now = 0L");
        for (ClassParameter parameter : list) {
            TypeVariableName t = TypeVariableName.get(parameter.classFullName);
            String resultName = parameter.classSimpleName + "_result";
            methodBuilder.addStatement("now = $T.currentTimeMillis()", System.class);
            methodBuilder.addStatement("$T " + parameter.classSimpleName + " = new $T()", t, t);
            methodBuilder.addStatement("$T " + resultName + " = new $T()", result, result);
            methodBuilder.addStatement(parameter.classSimpleName + ".initialize(context" + ", " + resultName + ")");
            methodBuilder.addStatement(resultName + ".useTime = System.currentTimeMillis() - now");
            methodBuilder.addStatement(resultName + ".name = \"" + parameter.starDriverInit.name() + "\"");
            methodBuilder.addStatement(resultName + ".className = \"" + parameter.classFullName + "\"");
            methodBuilder.addStatement("resultList.add(" + resultName + ")");
        }
        methodBuilder.addStatement("return resultList");
        return methodBuilder.build();
    }

    private static class ProcessClassNode implements Comparable<ProcessClassNode> {
        int count;
        ClassParameter value;
        List<ProcessClassNode> next = new ArrayList<>();

        public ProcessClassNode(ClassParameter value) {
            this.value = value;
        }

        @Override
        public int compareTo(ProcessClassNode processClassNode) {
            return count - processClassNode.count;
        }
    }

    private static class ClassParameter {
        Element element;
        StarDriverInit starDriverInit;

        String packageName;
        String simpleName;
        String classSimpleName;
        String classFullName;

        public ClassParameter(Elements elementUtils, Element element) {
            this.element = element;
            this.starDriverInit = element.getAnnotation(StarDriverInit.class);
            PackageElement packageElement = elementUtils.getPackageOf(this.element);
            this.packageName = packageElement.getQualifiedName().toString();
            this.classSimpleName = element.getSimpleName().toString();
            this.simpleName = this.starDriverInit.name();
            if (this.simpleName.isEmpty()) {
                this.simpleName = this.classSimpleName;
            }
            this.classFullName = element.asType().toString();
        }
    }
}
